%% On the Optimal Design of Transfers and Income-Tax Progressivity
% Consumption equivalents

clear; close all; clc; 
addpath('../../../../Matlab_Figure_Files');

%% Inputs (Results from Optimization)

% Parameters
sigma = 2.0;                 
beta = 0.9675;

% Number of grid points
ngrids = dlmread('params.txt');
na = ngrids(1);                 % number of asset grid points
nx = ngrids(2);                 % number of productivity grid points
ns = na*nx;                     % total number of idiosyncratic states

% Asset grid
a_vec = dlmread('a_vec.txt');
a_vec = reshape(a_vec',numel(a_vec),1); 
a_vec = reshape(a_vec(1:na),na,1);

% Transition matrix
px = dlmread('Px.txt');
px = px';
px = px(:);
px = px(1:nx*nx);
px = reshape(px,nx,nx);

% Asset policy initial steady state
a_pol_i = dlmread('a_pol_init.txt');
a_pol_i = reshape(a_pol_i',numel(a_pol_i),1);   
a_pol_i = reshape(a_pol_i(1:ns),ns,1);
a_pol_i = reshape(a_pol_i,na,nx);
a_pol_i = squeeze(a_pol_i);

% Labor policy initial steady state
h_pol_i = dlmread('h_pol_init.txt');
h_pol_i = reshape(h_pol_i',numel(h_pol_i),1);   
h_pol_i = reshape(h_pol_i(1:ns),ns,1);
h_pol_i = reshape(h_pol_i,na,nx);
h_pol_i = squeeze(h_pol_i);

% Consumption policy initial steady state
c_pol_i = dlmread('c_pol_init.txt');
c_pol_i = reshape(c_pol_i',numel(c_pol_i),1);   
c_pol_i = reshape(c_pol_i(1:ns),ns,1);
c_pol_i = reshape(c_pol_i,na,nx);
c_pol_i = squeeze(c_pol_i);

% Value funtion steady state
V_i = dlmread('V_init.txt');
V_i = reshape(V_i',numel(V_i),1);   
V_i = reshape(V_i(1:ns),ns,1);
V_i = reshape(V_i,na,nx);

% Measure initial steady state
mu_i = dlmread('mu_init.txt');
mu_i = reshape(mu_i',numel(mu_i),1);   
mu_i = reshape(mu_i(1:ns),ns,1);
mu_i = reshape(mu_i,na,nx);

% Value function transition
V_j = dlmread('V_TR.txt');
V_j = reshape(V_j',numel(V_j),1);   
V_j = reshape(V_j(1:ns),ns,1);
V_j = reshape(V_j,na,nx);

% Value funtion steady state
V_l = dlmread('V_lr.txt');
V_l = reshape(V_l',numel(V_l),1);   
V_l = reshape(V_l(1:ns),ns,1);
V_l = reshape(V_l,na,nx);

% Measure final steady state
mu_l = dlmread('mu_lr.txt');
mu_l = reshape(mu_l',numel(mu_l),1);   
mu_l = reshape(mu_l(1:ns),ns,1);
mu_l = reshape(mu_l,na,nx);


%% Compute Consumption equivalents

% Compute script C in the notes by iteration (C_aux here; 0 for the guess,1 as the update in the iteration)
C_aux_0 = zeros(na,nx);
C_aux_1 = zeros(na,nx);

for iter = 1:10000      % check non-binding max number of iterations
    for i = 1:na        % iterate over asset state
        for j = 1:nx    % iterate over productivity
            c_choice = c_pol_i(i,j);                    % consumption given state and policy function
            inst_ut = (c_choice^(1-sigma))/(1-sigma);   % instantaneous utility implied by consumption choice  
            C_aux_1(i,j) = inst_ut;                     % first component of C_aux 
            a_choice = a_pol_i(i,j);                    % asset choice given state and policy function
            a_low = max(sum(a_choice>a_vec),1);         % for linear interpolation: point on asset grid right below asset choice
            a_high = a_low + 1;                         % for linear interpolation: point on asset grid right above asset choice
            % Iterate over possible future productivity states
            for jp = 1:nx    
                % For each possible future productivity state: 
                % Compute the continuation value (v_aux) by linear interpolation given current guess for C_aux
                v_aux = C_aux_0(a_low,jp) + ((a_choice - a_vec(a_low))*(C_aux_0(a_high,jp) - C_aux_0(a_low,jp)))/(a_vec(a_high) - a_vec(a_low));
                % For each possible future productivity state:
                % Add the continuation value with discounting and the respective transition probability
                C_aux_1(i,j) = C_aux_1(i,j) + beta*px(j,jp)*v_aux;
            end
        end
    end
    if (max(max(abs(C_aux_1-C_aux_0))) < 1e-10)
        break
    else
        C_aux_0 = C_aux_1;
    end
end

% Formula for the consumption equivalent
gamma_mat = -1+(1+((V_j-V_i)./C_aux_1)).^(1/(1-sigma));

winners = gamma_mat > 0;
winners_share = sum(sum(winners.*mu_i));

% Average gain
ce_gain = sum(sum(sum(gamma_mat.*mu_i)));

%% Plot distribution of consumption equivalents by productivities

% Average consumption equivalents over productivities
ce_avg = (sum((gamma_mat.*mu_i))./sum(mu_i))';

% 20th and 80th percentile by productivity level
ce_20 = zeros(nx,1);
ce_80 = zeros(nx,1);
for i = 1:nx
    gamma_aux = gamma_mat(:,i);                 % CEs for all asset levels for this productivity
    mu_aux = mu_i(:,i);                         % corresponding measure
    mu_aux = mu_aux/sum(mu_aux);                % normalize this measure to one              
    [gamma_aux_sort,ind] = sort(gamma_aux);     % sort CEs for this productivity
    mu_aux_sort = mu_aux(ind);                  % sort measure accordingly
    mu_aux_cum = cumsum(mu_aux_sort);           % cumulative measure
    [~,ind] = min(abs(mu_aux_cum-0.20));        % location of 20th percentile
    ce_20(i) = gamma_aux_sort(ind);             % CE at 20th percentile
    [~,ind] = min(abs(mu_aux_cum-0.80));        % location of 80th percentile
    ce_80(i) = gamma_aux_sort(ind);             % CE at 80th percentile
end

% Stationary productivity distributiion
level_mix = dlmread('x_vec.txt')';              % productivity states
level_mix = level_mix(1:19)';                   % productivity states
dist_mix = ones(19,1)/19;                       % allocate stationary distribution
dist_mix_aux = dist_mix;                        % allocate stationary distribution
err = 1;                                        % initialize error
while err > 1e-10
    dist_mix_aux = (dist_mix'*px)';             % update distribution with transition matrix
    err = max(abs(dist_mix-dist_mix_aux));      % error
    dist_mix = dist_mix_aux;                    % update guess
end

% Variables for plot
x_aux = [level_mix',fliplr(level_mix')];
inBetween = [ce_20'*100,fliplr(ce_80')*100];

% Plot
fig = figure(101);
yyaxis left
h0 = plot(level_mix,ce_20*100,'--',level_mix,ce_80*100,'--','Color',[0.85, 0.9170, 0.9910],'Linewidth',2); hold on
fill(x_aux, inBetween, [0.85, 0.9170, 0.9910]); hold on
h1 = plot(level_mix,ce_avg*100,'-','Color',[0 0.4470 0.7410],'LineWidth',2); hold on
ylabel('CE (\%)','Interpreter','LaTex','Fontsize',12)
set(gca,'XGrid','off','YGrid','on','Fontsize',12) 
ylim([-10 30])
yyaxis right
h2 = plot(level_mix,100*dist_mix,':','Color',[0.8500 0.3250 0.0980],'LineWidth',2);
leg = legend([h1 h2],{'Consumption equivalent','Measure'});
set(leg,'Interpreter','LaTex','Location','NorthEast','Fontsize',12)
legend boxoff
ylabel('Measure (\%)','Interpreter','LaTex','Fontsize',12)
xlabel('Productivity','Interpreter','LaTex','Fontsize',12)
set(gca,'XGrid','off','YGrid','on','Fontsize',12) 
xlim([0 7])
set(gca,'XGrid','off','YGrid','on','Fontsize',12) 
set(gca,'TickLabelInterpreter','LaTex')
ylim([-5 15])
yline(0,':','Color','black','Linewidth',2,'HandleVisibility','off') 
run graph_extended.m
saveas(gcf,'qm_CE_prods.pdf')

fig = figure(102);
yyaxis left
h0 = plot(level_mix,ce_20*100,'--',level_mix,ce_80*100,'--','Color',[0.8 0.8 0.8],'Linewidth',2); hold on
fill(x_aux, inBetween, [0.8 0.8 0.8]); hold on
h1 = plot(level_mix,ce_avg*100,'-','Color',[0.4 0.4 0.4],'LineWidth',2); hold on
ylabel('CE (\%)','Interpreter','LaTex','Fontsize',12)
set(gca,'XGrid','off','YGrid','on','Fontsize',12) 
ylim([-10 30])
yyaxis right
h2 = plot(level_mix,100*dist_mix,':','Color',[0 0 0],'LineWidth',2);
leg = legend([h1 h2],{'Consumption equivalent','Measure'});
set(leg,'Interpreter','LaTex','Location','NorthEast','Fontsize',12)
legend boxoff
ylabel('Measure (\%)','Interpreter','LaTex','Fontsize',12)
xlabel('Productivity','Interpreter','LaTex','Fontsize',12)
set(gca,'XGrid','off','YGrid','on','Fontsize',12) 
xlim([0 7])
set(gca,'XGrid','off','YGrid','on','Fontsize',12) 
set(gca,'TickLabelInterpreter','LaTex')
ylim([-5 15])
yline(0,':','Color','black','Linewidth',2,'HandleVisibility','off') 
ax = gca;
ax.YAxis(1).Color = '[0.4 0.4 0.4]';
ax.YAxis(2).Color = '[0 0 0]';
run graph_extended.m
saveas(gcf,'qm_CE_prods_bw.eps')

